import random

import gym
from gym.spaces import Discrete

from q_table.Games.GameMaze import GameMaze
from q_table.Games.MazeDTO import generate_maze


class MazeEnv(gym.Env):
    def __init__(self, minsize=2, maxsize=2):
        self.gameMaze = GameMaze(minsize, maxsize)
        self.size = self.gameMaze.random_maze_size()
        self.Tmaze = generate_maze(self.size, self.size)
        self.Maze = self.Tmaze[0]
        self.maze = self.Maze.cellMaze
        start = self.Tmaze[1]
        end = self.Tmaze[2]
        self.start_cell = self.maze[start[0]][start[1]]
        self.end_cell = self.maze[end[0]][end[1]]
        self.end_cell.is_end = True
        self.cur_pos = self.start_cell
        self.dirs = [[-1, 0], [1, 0], [0, 1], [0, -1]]
        self.action_space = Discrete(4)
        self.observation_space = Discrete(self.size * self.size)

    def reset(self, *args, **kwargs):
        maze = [[1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 1, 0, 0, 0, 1], [1, 1, 1, 0, 1, 0, 1, 0, 1],
                [1, 0, 0, 0, 0, 0, 1, 0, 1], [1, 1, 1, 1, 1, 1, 1, 0, 1], [1, 0, 0, 0, 1, 0, 0, 0, 1],
                [1, 0, 1, 1, 1, 1, 1, 0, 1], [1, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1]]

        self.Maze.reset_cellMaze(maze)
        self.maze = self.Maze.cellMaze
        self.cur_pos = self.start_cell
        loc = self.cur_pos.loc
        return loc[0] * self.size + loc[1]

    def step(self, action, render=False):
        dir = self.dirs[action]
        loc_pre = self.cur_pos.loc
        state_pre = loc_pre[0] * self.size + loc_pre[1]
        self.cur_pos = self.gameMaze.move(self.maze, self.cur_pos, self.end_cell, dir, render)
        loc = self.cur_pos.loc
        print(f'loc_pre: {loc_pre}, loc: {loc}')
        state = loc[0] * self.size + loc[1]
        done = self.end_cell.loc == self.cur_pos.loc
        if done:
            reward = self.size * self.size * self.size * self.size
        elif state_pre == state:
            reward = -(self.size ** 2)
        else:
            reward = -(self.cur_pos.visited_num + 1)
            # reward = -1
        # time.sleep(0.1)
        print(f'reward: {reward}')
        return state, reward, done, {}


if __name__ == '__main__':
    env = MazeEnv(maxsize=3, minsize=3)
    env.reset()
    for i in range(100):
        step = random.randint(0, 3)
        env.step(step)
